# -*- coding: utf-8 -*-
import os
from tools.writing_classes import writing_classes
from tools.checkJpgXml import checkJX
from tools.transform_labels import transform_lables

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
   ##########设置参数
   #####输入数据集路径
   images_dir = '/home/**/dataset/traffic_signs/VOCTRY/JPEGImages'
   annotations_dir = '/home/**/dataset/traffic_signs/1080P/VOC1080p/Annotations'
   classes = ['pl120', 'w57', 'w55', 'pl100', 'ip', 'w59', 'w13', 'pl5', 'pl80', 'p27', 'p26', 'p23',
              'ph5', 'ph4', 'pl20', 'pl50', 'il100', 'pl60', 'pl40', 'pg', 'pm20', 'pn', 'il80', 'p3',
              'i2', 'i5', 'i4', 'p10', 'p11', 'p12', 'p19', 'pne', 'ph4.5', 'pl30', 'p6', 'p5',
              'il60', 'w30', 'w32','pm55','pm30','pl70']
   index_txt_path = '/home/**/dataset/traffic_signs/VOCTRY/ImageSets/Main/'
   test_txt_path = '/home/**/dataset/traffic_signs/VOCTRY/ImageSets/Main/test.txt'


   #####输出数据集路径
   yolo_dir = '/home/**/dataset/yolov3'
   ylo_images_dir = '/home/**/dataset/yolov3/images'

   ###1. 输出classes.names文件
   print("=" * 6, "\t1. outputing a file named classes.names\t", "=" * 6)
  # writing_classes(classes, yolo_dir)

   ###2.检查jpg和xml文件是否是一一对应的
   print("=" * 6, "\t2. checking jpg and xml\t", "=" * 6)
   checkJX(jpeg_dir=images_dir,anno_dir=annotations_dir)

   ###3.读取ImageSets/Main/文件夹中的train.txt和test.txt，换成需要的格式.然后构建pytorch版本yolov3格式。
   print("=" * 6, "\t3. outting the index files: train.txt and test.txt \t", "=" * 6)
   transform_lables(images_dir,yolo_dir,index_txt_path,classes)





